#define _USE_MATH_DEFINES
#include "PathOptProb.h"
#include "DebugConfig.h"
#include <sstream>
#include <stdexcept>
#include <limits>
#include <cmath>     // For std::round, std::sin, std::cos, etc.
#include <iostream>  // For std::cout
#include <chrono>    // For timing
#include <random>    // For std::mt19937 and std::uniform_real_distribution

namespace pagmo_problems
{
    PathOptProb::PathOptProb(algo::AlgoEvaluator* evaluator,
                             const algobase::TaskConfig* config,
                             const algobase::EvaluateL1IndexMap* weights,
                             size_t num_control_points,
                             algo::AlgoEvaluatorDB* db)
        : m_evaluator(evaluator),
          m_config(config),
          m_weights(weights),
          m_num_control_points(num_control_points),
          m_db(db)
    {
        if (!m_evaluator || !m_config || !m_weights || !m_db) {
            throw std::invalid_argument("Evaluator, config, weights, and db must not be null.");
        }
        if (m_num_control_points == 0) {
            throw std::invalid_argument("Number of control points must be greater than 0.");
        }
    }

    std::vector<double> PathOptProb::fitness(const std::vector<double>& dv) const
    {
        static int fitness_call_count = 0;
        fitness_call_count++;
        std::cout << "\n--- [PathOptProb] Fitness call #" << fitness_call_count << " started ---" << std::endl;

        algobase::Path path = decode(dv);
        double total_result = 0.0;
        algobase::EvaluatorCoords result_coords;

        std::cout << "  [PathOptProb] Calling evaluator->exec()..." << std::endl;
        auto start_time = std::chrono::high_resolution_clock::now();
        std::map<algobase::EvaluateL2IndexType, double> L2weight;
        int eval_status = m_evaluator->exec(path, total_result, result_coords);

        auto end_time = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
        std::cout << "  [PathOptProb] Evaluator->exec() finished. Status: " << eval_status 
                  << ". Duration: " << duration.count() << " ms." << std::endl;

        if (eval_status != 0 || result_coords.empty()) {
            // Return a penalty for each objective
            return std::vector<double>(get_nobj(), std::numeric_limits<double>::max());
        }

        // Aggregate L1 evaluation indices over the path (e.g., by averaging)
        std::vector<double> avg_l1_values(get_nobj(), 0.0);
        for (const auto& coord : result_coords) {
            for (size_t i = 0; i < get_nobj(); ++i) {
                avg_l1_values[i] += coord.mL1EvaIndexArray[i];
            }
        }

        for (size_t i = 0; i < get_nobj(); ++i) {
            avg_l1_values[i] /= result_coords.size();
            // Assuming higher scores are better, we negate them for minimization.
            avg_l1_values[i] = -avg_l1_values[i];
        }

        std::stringstream ss;
        ss << "Path fitness calculated. Objectives: ";
        for(size_t i = 0; i < avg_l1_values.size(); ++i) {
            ss << avg_l1_values[i] << (i == avg_l1_values.size() - 1 ? "" : ", ");
        }
        DEBUG_LOG("AlgoNSGA.log_fitness_calculation", ss.str());
        std::cout << "  [PathOptProb] " << ss.str() << std::endl;

        return avg_l1_values;
    }

    std::pair<std::vector<double>, std::vector<double>> PathOptProb::get_bounds() const
    {
        size_t dim = m_num_control_points * 3; // lon, lat, depth for each control point
        std::vector<double> lb(dim);
        std::vector<double> ub(dim);

        const double min_clearance = 20.0; // meters above seabed
        const double max_depth_abs = 500.0; // max absolute depth
        const double surface_depth = -5.0; // meters (just below surface)

        for (size_t i = 0; i < m_num_control_points; ++i) {
            // Interpolate lon/lat for this control point to query terrain
            double ratio = static_cast<double>(i + 1) / (m_num_control_points + 1);
            double lon = m_config->mStartPos[0] * (1.0 - ratio) + m_config->mEndPos[0] * ratio;
            double lat = m_config->mStartPos[1] * (1.0 - ratio) + m_config->mEndPos[1] * ratio;

            // Longitude bounds
            lb[i * 3] = m_config->mMinLongitude;
            ub[i * 3] = m_config->mMaxLongitude;
            // Latitude bounds
            lb[i * 3 + 1] = m_config->mMinLatitude;
            ub[i * 3 + 1] = m_config->mMaxLatitude;
            
            // Depth bounds based on terrain
            algobase::CoordZ query_coord = {lon, lat, -100.0}; // depth is irrelevant for terrain query
            double terrain_altitude = m_db->getData(algobase::EEDT_Depth, m_config->mStartTime, query_coord, 9999.0f);
            
            double lower_bound;
            double upper_bound = surface_depth;

            if (terrain_altitude == 9999.0f) {
                // If terrain data is unavailable, use a conservative general depth range
                lower_bound = -200.0;
            } else {
                lower_bound = terrain_altitude + min_clearance;
            }

            lb[i * 3 + 2] = std::min(lower_bound, upper_bound); // lower_bound is more negative
            ub[i * 3 + 2] = std::max(lower_bound, upper_bound);
        }
        return {lb, ub};
    }

    size_t PathOptProb::get_nobj() const
    {
        // We have 4 objectives: Safety, Stealth, Economy, Completion.
        // EL1IT_Count is the size of the enum, so the number of objectives is one less.
        return algobase::EL1IT_Count - 1;
    }

    algobase::Path PathOptProb::decode(const std::vector<double>& dv) const
    {
        algobase::Path path;
        
        algobase::PathCoord start_point;
        start_point.mCoordZ[0] = m_config->mStartPos[0];
        start_point.mCoordZ[1] = m_config->mStartPos[1];
        start_point.mCoordZ[2] = -100.0; // Start depth is fixed
        path.push_back(start_point);

        for (size_t i = 0; i < m_num_control_points; ++i) {
            algobase::PathCoord control_point;
            control_point.mCoordZ[0] = dv[i * 3];
            control_point.mCoordZ[1] = dv[i * 3 + 1];
            control_point.mCoordZ[2] = dv[i * 3 + 2];
            path.push_back(control_point);
        }

        algobase::PathCoord end_point;
        end_point.mCoordZ[0] = m_config->mEndPos[0];
        end_point.mCoordZ[1] = m_config->mEndPos[1];
        end_point.mCoordZ[2] = -100.0; // End depth is fixed
        path.push_back(end_point);

        // --- Post-process to set Speed and Heading ---
        const double constant_speed_mps = 5.0; // meters per second

        for (size_t i = 0; i < path.size() - 1; ++i) {
            path[i].mSpeed = constant_speed_mps;
            
            double lon1 = path[i].mCoordZ[0];
            double lat1 = path[i].mCoordZ[1];
            double lon2 = path[i+1].mCoordZ[0];
            double lat2 = path[i+1].mCoordZ[1];

            // Convert to radians for calculation
            double lat1_rad = lat1 * M_PI / 180.0;
            double lat2_rad = lat2 * M_PI / 180.0;
            double dLon_rad = (lon2 - lon1) * M_PI / 180.0;

            double y = std::sin(dLon_rad) * std::cos(lat2_rad);
            double x = std::cos(lat1_rad) * std::sin(lat2_rad) - std::sin(lat1_rad) * std::cos(lat2_rad) * std::cos(dLon_rad);
            double heading_rad = std::atan2(y, x);
            
            path[i].mHeading = fmod((heading_rad * 180.0 / M_PI + 360.0), 360.0); // Convert to degrees
        }
        // Last point has no next point, so speed is 0 and heading is same as previous.
        path.back().mSpeed = 0.0;
        if (path.size() > 1) {
            path.back().mHeading = path[path.size() - 2].mHeading;
        }

        return path;
    }

    std::vector<double> PathOptProb::encode(const algobase::Path& path) const
    {
        std::vector<double> dv;
        dv.reserve(m_num_control_points * 3);

        if (path.size() <= 2) {
            // If path is just start and end, linearly interpolate control points
            for (size_t i = 0; i < m_num_control_points; ++i) {
                double ratio = static_cast<double>(i + 1) / (m_num_control_points + 1);
                dv.push_back(m_config->mStartPos[0] * (1.0 - ratio) + m_config->mEndPos[0] * ratio);
                dv.push_back(m_config->mStartPos[1] * (1.0 - ratio) + m_config->mEndPos[1] * ratio);
                dv.push_back(-100.0); // Default depth
            }
            return dv;
        }

        const size_t num_interior_points = path.size() - 2;
        for (size_t i = 0; i < m_num_control_points; ++i) {
            double sample_index_double = 1.0 + static_cast<double>(i * num_interior_points) / m_num_control_points;
            size_t sample_index = static_cast<size_t>(std::round(sample_index_double));
            sample_index = std::max(static_cast<size_t>(1), std::min(path.size() - 2, sample_index));
            
            dv.push_back(path[sample_index].mCoordZ[0]);
            dv.push_back(path[sample_index].mCoordZ[1]);
            dv.push_back(path[sample_index].mCoordZ[2]);
        }

        return dv;
    }

    algobase::Path PathOptProb::mutate(const algobase::Path& path, double mutation_strength) const
    {
        algobase::Path mutated_path = path;
        auto bounds = get_bounds();
        const auto& lb = bounds.first;
        const auto& ub = bounds.second;

        // Use C++ standard library for random number generation
        static std::random_device rd;
        static std::mt19937 gen(rd());

        // Mutate interior control points
        if (mutated_path.size() > 2) {
            for (size_t i = 1; i < mutated_path.size() - 1; ++i) {
                // Mutate Longitude
                double lon_range = (ub[ (i-1) * 3 ] - lb[ (i-1) * 3 ]) * mutation_strength;
                std::uniform_real_distribution<> lon_dist(-lon_range, lon_range);
                mutated_path[i].mCoordZ[0] += lon_dist(gen);
                mutated_path[i].mCoordZ[0] = std::max(lb[ (i-1) * 3 ], std::min(ub[ (i-1) * 3 ], mutated_path[i].mCoordZ[0]));

                // Mutate Latitude
                double lat_range = (ub[ (i-1) * 3 + 1 ] - lb[ (i-1) * 3 + 1 ]) * mutation_strength;
                std::uniform_real_distribution<> lat_dist(-lat_range, lat_range);
                mutated_path[i].mCoordZ[1] += lat_dist(gen);
                mutated_path[i].mCoordZ[1] = std::max(lb[ (i-1) * 3 + 1 ], std::min(ub[ (i-1) * 3 + 1 ], mutated_path[i].mCoordZ[1]));

                // Mutate Depth
                double depth_range = (ub[ (i-1) * 3 + 2 ] - lb[ (i-1) * 3 + 2 ]) * mutation_strength;
                std::uniform_real_distribution<> depth_dist(-depth_range, depth_range);
                mutated_path[i].mCoordZ[2] += depth_dist(gen);
                mutated_path[i].mCoordZ[2] = std::max(lb[ (i-1) * 3 + 2 ], std::min(ub[ (i-1) * 3 + 2 ], mutated_path[i].mCoordZ[2]));
            }
        }
        return mutated_path;
    }
}